#/*************************************************************************
#    > File Name: ROOT_Operation.py
#    > Author: Yan Wang
#    > Mail: wangyan@imnu.edu.cn
#    > Created Time: Wed 24 Aug 2022 11:39:01 PM CST
# ************************************************************************/
#!/usr/bin/env python
# -*- coding: utf-8 -*-

#import ROOT
from Library import *
##############################################################################
# 导入from delphes
#delphes_dir = MG_para.Param_Dir["MG"] + "Delphes"
#sys.path.append(delphes_dir)
Delphes_Dir="/home/yancywww/Software/bin/MG5_aMC_v3_5_3/Delphes"
ROOT.gSystem.Load(Delphes_Dir+"/libDelphes.so")
#ROOT.gSystem.Load("libDelphes")

def ROOT_Running_Code(root_file_name,commands=[]):
    # 判断是否存在该root文件
    if (ROOT.gSystem.AccessPathName(root_file_name)) :
        ROOT.Info("example1.py", root_file_name+" does not exist")
        exit()

# 读取树干Delphes的root trees
    chain = ROOT.TChain("Delphes")
    chain.Add(root_file_name) # 读取文件
# Create object of class ExRootTreeReader
    treeReader = ROOT.ExRootTreeReader(chain)
    results=[]
    for com in commands:
        #res=os.system(com(treeReader))
        res=com(treeReader)
        results.append(res)
    return(results)

def ROOT_Get_AllEntries_Number(file_name):
    file_type=type(file_name)
    if file_type is str:
        number_of_Entries = ROOT_Get_AllEntries_Number_file(file_name)
    elif file_type is pathlib.PosixPath:
        file_str = str(file_name)
        number_of_Entries = ROOT_Get_AllEntries_Number_file(file_str)
    else :
        number_of_Entries = ROOT_Get_AllEntries_Number_tree(file_name)
#    else:
#        print("wrong input file type, oly str or treereader is acceptable, current is:", type(file_name))

        #file_type is cppyy.gbl.ExRootTreeReader:
    return(number_of_Entries)



def ROOT_Get_AllEntries_Number_file(root_file_name):
    # 判断是否存在该root文件
    if (ROOT.gSystem.AccessPathName(root_file_name)) :
        ROOT.Info("example1.py", root_file_name+" does not exist")
        exit()

# 读取树干Delphes的root trees
    chain = ROOT.TChain("Delphes")
    chain.Add(root_file_name) # 读取文件
# Create object of class ExRootTreeReader
    treeReader = ROOT.ExRootTreeReader(chain)
    number_of_Entries = treeReader.GetEntries() # 获得事件数
    return(number_of_Entries)


def ROOT_Get_AllEntries_Number_tree(treeReader):
    numberOfEntries = treeReader.GetEntries() # 获得事件数
    print("The total event number is:", numberOfEntries)
    return(numberOfEntries)

def ROOT_Get_FinalState_Number(treeReader):
    num_particle=np.zeros(5)
    num_elec    =np.zeros(5)
    num_muon    =np.zeros(5)
    num_lep     =np.zeros(5)
    num_jet     =np.zeros(5)
    num_bjet    =np.zeros(5)
    all_num={"particle":num_particle, "elec":num_elec, "muon":num_muon, "lep":num_lep, "jet":num_jet, "bjet":num_bjet}

    branchParticle=treeReader.UseBranch("Particle")
    branchElectron=treeReader.UseBranch("Electron")
    branchMuon    =treeReader.UseBranch("Muon")
    branchJet     =treeReader.UseBranch("Jet")

    for entry in range(treeReader.GetEntries()):
        treeReader.ReadEntry(entry)
        nbj=0
        for nj in range(branchJet.GetEntriesFast()):
            jet=branchJet.At(nj)
            if jet.BTag == True:
                nbj+=1
        for num in range(5):
            if branchParticle.GetEntriesFast() == num:
                num_particle[num]+=1
            if branchElectron.GetEntriesFast() == num:
                num_elec[num]+=1
            if branchMuon.GetEntriesFast() == num:
                num_muon[num]+=1
            if branchElectron.GetEntriesFast()+branchMuon.GetEntriesFast() == num:
                num_lep[num]+=1
            if branchJet.GetEntriesFast() == num:
                num_jet[num]+=1
            if nbj == num:
                num_bjet[num]+=1

    print("         numbers are : " , ["0,   1,   2,   3,   4"])
    print("Particle numbers are : " , num_particle)
    print("Electron numbers are : " , num_elec    )
    print("Muon     numbers are : " , num_muon    )
    print("Lepton   numbers are : " , num_lep)
    print("Jet      numbers are : " , num_jet     )
    print("BJet     numbers are : " , num_bjet    )
    return(all_num)


